#include "SimpleScreenCaptureHandler.h"
#include <iostream>

SimpleScreenCaptureHandler::SimpleScreenCaptureHandler(osgViewer::Viewer* viewer, const std::string& saveDir)
    : _viewer(viewer), _saveDir(saveDir), _captureNextFrame(false)
{
}

bool SimpleScreenCaptureHandler::handle(const osgGA::GUIEventAdapter& ea, osgGA::GUIActionAdapter&)
{
    if (ea.getEventType() == osgGA::GUIEventAdapter::KEYDOWN)
    {
        if (ea.getKey() == 'c' || ea.getKey() == 'C')
        {
            _captureNextFrame = true;
            return true;
        }
    }
    return false;
}

void SimpleScreenCaptureHandler::capture()
{
    if (!_captureNextFrame)
        return;

    osg::Camera* camera = _viewer->getCamera();
    int width = camera->getViewport()->width();
    int height = camera->getViewport()->height();

    osg::ref_ptr<osg::Image> image = new osg::Image();
    image->allocateImage(width, height, 1, GL_RGBA, GL_UNSIGNED_BYTE);

    camera->attach(osg::Camera::COLOR_BUFFER, image.get());

    _viewer->frame();  // 渲染一帧用于填充缓冲

    camera->detach(osg::Camera::COLOR_BUFFER);

    std::string filename = _saveDir + "/screenshot.png";

    if (osgDB::writeImageFile(*image, filename))
        std::cout << "截图成功: " << filename << std::endl;
    else
        std::cout << "截图失败" << std::endl;

    _captureNextFrame = false;
}
